from .base_test import BaseTest


class CorrelationTest(BaseTest):

	def test_use_index_only_scan(self):
		node = self.node
		node.start()
		node.safe_psql("""
			CREATE EXTENSION orioledb;
		""")

		node.safe_psql("""
			CREATE TABLE pg_foobar (
				id SERIAL,
				k INTEGER DEFAULT '0' NOT NULL,
				c CHAR(120) DEFAULT '' NOT NULL,
				pad CHAR(60) DEFAULT '' NOT NULL,
				PRIMARY KEY (id)
			);

			insert into pg_foobar (id, k) select a, 9000000 + a from generate_series(1, 100000) a;
			CREATE INDEX pg_foobar_ix ON pg_foobar(k);
			analyze verbose pg_foobar;
		""")

		node.safe_psql("""
			CREATE TABLE o_foobar (
				id SERIAL,
				k INTEGER DEFAULT '0' NOT NULL,
				c CHAR(120) DEFAULT '' NOT NULL,
				pad CHAR(60) DEFAULT '' NOT NULL,
				PRIMARY KEY (id)
			) USING orioledb;

			insert into o_foobar (id, k) select a, 9000000 + a from generate_series(1, 100000) a;
			CREATE INDEX o_foobar_ix ON o_foobar(k);
			analyze verbose o_foobar;
		""")

		correlations = node.execute("""
			SELECT tablename, attname, correlation FROM pg_stats
				WHERE tablename IN ('pg_foobar', 'o_foobar')
				ORDER BY attname, tablename;
		""")
		self.assertEqual(correlations, [('o_foobar', 'c', 1.0),
		                                ('pg_foobar', 'c', 1.0),
		                                ('o_foobar', 'id', 1.0),
		                                ('pg_foobar', 'id', 1.0),
		                                ('o_foobar', 'k', 1.0),
		                                ('pg_foobar', 'k', 1.0),
		                                ('o_foobar', 'pad', 1.0),
		                                ('pg_foobar', 'pad', 1.0)])

		rows = """
			5642924, 23728, 8070388, 8165335, 9013596,
			6796988, 2757245, 779777, 5393690, 2174111,
			7288676, 9203419, 1164884,   2858736, 6597995,
			2457392, 7016879, 366253, 7009025, 7598636, 3651827,
			4511813, 6140952, 109135, 9585912, 4241228, 6085362,
			3610730, 7218735, 5008924, 4459230, 4778645,
			8786501, 2162800, 5975982, 2424563, 3929852,
			5839423, 4537270, 8314863, 2222310, 1228862,
			6678420, 2000071, 6453377, 4871873, 1421519,
			7956413, 4206990, 2059264, 687992, 3524116, 7587596,
			2563436, 5472222, 3180146, 6288357, 2819706, 7413802,
			9756837, 5597154, 1589924, 1609178, 6743605, 5447307,
			8583654, 2728526, 661395, 9224073, 6797236, 2605946,
			6906407, 1702528, 2005661, 1076111, 8109872, 87456,
			3809676, 5498674, 6270243, 8508471, 532157, 5636831,
			9401488, 8614377, 4323511, 7097436, 4997189, 2808874,
			3356750, 7384619, 386615, 238166, 7559385, 1069848,
			731617, 7404044, 3982062, 3312563, 5075720, 6671700 ,
			1717800 ,5173800 ,691200 ,1489200 ,4254000 ,4927800 ,
			4012500 ,1756200 ,7003800 ,7743300 ,6614100 ,3469800 ,
			2670600 ,6720000 ,1082400 ,6293400 ,4722000 ,5057700 ,
			8149200 ,7025100 ,4149300 ,7247100 ,2460600 ,7757700 ,
			1879200 ,2595300 ,7648800 ,3123600 ,61200 ,3231000 ,
			7534500 ,8771100 ,5860200 ,8754900 ,9114600 ,7832700 ,
			314700 ,1882500 ,4216500 ,6017100 ,5186100 ,9347400 ,
			3356700 ,2501400 ,8491800 ,363600 ,8633700 ,2027400 ,
			5301600 ,4505400 ,6594000 ,3534000 ,7902600 ,4821600 ,
			73500 ,1265400 ,5233500 ,837600 ,8284500 ,3709800 ,
			6751200 ,8503200 ,8467500 ,7413300 ,4329900 ,1093800 ,
			362100 ,4272600 ,5555700 ,1104300 ,3238800 ,2871900 ,
			961200 ,8075700 ,62400 ,1880400 ,2653800 ,5329800 ,
			8865000 ,2763900 ,7145700 ,7699800 ,4867800 ,794700 ,
			3978000 ,2410800 ,1632900 ,5543400 ,2366100 ,5581500 ,
			1892700 ,3283500 ,6501300 ,8085600 ,4548900 ,687900 ,
			4792500 ,703800 ,7394700 , 2068500 ,3705300 ,7523100 ,
			7069500 ,5616000 ,5248200 ,8563200 ,1912500 ,8197200 ,
			2339700 ,1540500 ,7773900 ,5489700 ,7511400 ,7870200 ,
			1238100 ,3084600 ,8630100 ,4998000 ,9706800 ,1254900 ,
			7380600 ,1870200 ,2190000 ,2933100 ,261600 ,4392600 ,
			5787000 ,5311800 ,756300 ,7248000 ,4387800 ,2080500 ,
			5779200 ,8440500 ,4569900 ,1925700 ,2067300 ,591300 ,
			9026100 ,9273300 ,7678800 ,3800400 ,6113400 ,441900 ,
			2657700 ,7502100 ,7443300 ,784800 ,6900900 ,8064000 ,
			8601600 ,6960600 ,2658600 ,6205800 ,648900 ,6055800 ,
			3068400 ,6144000 ,453900 ,9291300 ,9025200 ,5915400 ,
			8305500 ,7938300 ,524100 ,2224200 ,2900700 ,5844600 ,
			2503500 ,6321900 ,8624700 ,2529600 ,2625000 ,5051400 ,
			8235900 ,9353400 ,5421900 ,2052900 ,8472900 ,2511000 ,
			7499700 ,57900 ,389700 ,2400600 ,4624200 ,176700 ,7220700 ,
			8481300 ,1951500 ,9264000 ,6066300 ,4145400 ,3258900 ,6990600 ,
			4667400 ,7416000 ,2022900 ,6566700 ,6832200 , 1222800 ,5676900 ,
			2955000 ,443700 ,6814500 ,2179500 ,6851700 ,8048700 ,6593700 ,
			3240000 ,7113600 ,4159200 ,3660900 ,6333600 ,7261500 ,3186300 ,
			9779100 ,2054400 ,2572800 ,4650300 ,5175000 ,5133000 ,1043400 ,
			4808100 ,9529800 ,2898900 ,2839800 ,6906900 ,5349000 ,7386900 ,
			485100 ,672000 ,1880100 , 3112200 ,4277700 ,9450300 ,5631900 ,
			3878400 ,3218700 ,9110100 ,2155200 ,8022300 ,4415100 ,2084100 ,
			7807800 ,7320000 ,5898900 ,1249500 ,8752800 ,1992900 ,2874900 ,
			578100 ,2159400 ,39000 ,9630900 ,733200 ,5080200 ,8530200 ,418200 ,
			1581000 ,2121600 ,124200 ,3182100 ,7698900 ,6715200 ,5477700 ,3286500 ,
			8903400 ,7191900 ,4154400 ,7607400 ,4493400 ,8313000 ,2258100 ,581100 ,
			2082900 ,6379500 ,7670400 ,824700 ,6763500 ,4615500 ,5935200 ,155400 ,
			7566900 ,4012500 ,6375300 ,498300 ,4818000 ,8429700 ,9403500 ,6930000 ,
			3385500 ,9165600 ,8497500 ,4278000 ,1045800 ,8556600 ,5516400 ,4043700 ,4906800
		"""
		plan = node.execute(f"""
			EXPLAIN (COSTS OFF, FORMAT JSON)
				SELECT k FROM pg_foobar WHERE k IN ({rows});
		""")[0][0][0]["Plan"]
		self.assertEqual('Index Only Scan', plan["Node Type"])
		self.assertEqual('pg_foobar_ix', plan['Index Name'])

		plan = node.execute(f"""
			EXPLAIN (COSTS OFF, FORMAT JSON)
				SELECT k FROM o_foobar WHERE k IN ({rows});
		""")[0][0][0]["Plan"]
		self.assertEqual('Index Only Scan', plan["Node Type"])
		self.assertEqual('o_foobar_ix', plan['Index Name'])
